// addrspace.cc 
//	Routines to manage address spaces (executing user programs).
//
//	In order to run a user program, you must:
//
//	1. link with the -N -T 0 option 
//	2. run coff2noff to convert the object file to Nachos format
//		(Nachos object code format is essentially just a simpler
//		version of the UNIX executable object code format)
//	3. load the NOFF file into the Nachos file system
//		(if you haven't implemented the file system yet, you
//		don't need to do this last step)
//
// Copyright (c) 1992-1993 The Regents of the University of California.
// All rights reserved.  See copyright.h for copyright notice and limitation 
// of liability and disclaimer of warranty provisions.

#include "copyright.h"
#include "system.h"
#include "addrspace.h"
#include "noff.h"
#include "table.h"
#include "synch.h"

extern "C" { int bzero(char *, int); };

Table::Table(int s) : map(s), table(0), lock(0), size(s) {
    table = new void *[size];
    lock = new Lock("TableLock");
}

Table::~Table() {
    if (table) {
	delete table;
	table = 0;
    }
    if (lock) {
	delete lock;
	lock = 0;
    }
}

void *Table::Get(int i) {
    // Return the element associated with the given if, or 0 if
    // there is none.

    return (i >=0 && i < size && map.Test(i)) ? table[i] : 0;
}

int Table::Put(void *f) {
    // Put the element in the table and return the slot it used.  Use a
    // lock so 2 files don't get the same space.
    int i;	// to find the next slot

    lock->Acquire();
    i = map.Find();
    lock->Release();
    if ( i != -1) {
		table[i] = f;
		itemNum ++;
	}
    return i;
}

void *Table::Remove(int i) {
    // Remove the element associated with identifier i from the table,
    // and return it.

    void *f =0;

    if ( i >= 0 && i < size ) {
	lock->Acquire();
	if ( map.Test(i) ) {
	    map.Clear(i);
	    f = table[i];
	    table[i] = 0;
		itemNum --;
	}
	lock->Release();
    }
    return f;
}

//----------------------------------------------------------------------
// SwapHeader
// 	Do little endian to big endian conversion on the bytes in the 
//	object file header, in case the file was generated on a little
//	endian machine, and we're now running on a big endian machine.
//----------------------------------------------------------------------

static void 
SwapHeader (NoffHeader *noffH)
{
	noffH->noffMagic = WordToHost(noffH->noffMagic);
	noffH->code.size = WordToHost(noffH->code.size);
	noffH->code.virtualAddr = WordToHost(noffH->code.virtualAddr);
	noffH->code.inFileAddr = WordToHost(noffH->code.inFileAddr);
	noffH->initData.size = WordToHost(noffH->initData.size);
	noffH->initData.virtualAddr = WordToHost(noffH->initData.virtualAddr);
	noffH->initData.inFileAddr = WordToHost(noffH->initData.inFileAddr);
	noffH->uninitData.size = WordToHost(noffH->uninitData.size);
	noffH->uninitData.virtualAddr = WordToHost(noffH->uninitData.virtualAddr);
	noffH->uninitData.inFileAddr = WordToHost(noffH->uninitData.inFileAddr);
}

//----------------------------------------------------------------------
// AddrSpace::AddrSpace
// 	Create an address space to run a user program.
//	Load the program from a file "executable", and set everything
//	up so that we can start executing user instructions.
//
//	Assumes that the object code file is in NOFF format.
//
//	"executable" is the file containing the object code to load into memory
//
//      It's possible to fail to fully construct the address space for
//      several reasons, including being unable to allocate memory,
//      and being unable to read key parts of the executable.
//      Incompletely consretucted address spaces have the member
//      constructed set to false.
//----------------------------------------------------------------------

AddrSpace::AddrSpace(OpenFile *executable) : fileTable(MaxOpenFiles) {
    NoffHeader noffH;
    unsigned int i, size, codeDataSize;
	
    // Don't allocate the input or output to disk files
    fileTable.Put(0);
    fileTable.Put(0);

    execFile = executable;	// Keep the executable file
    executable->ReadAt((char *)&noffH, sizeof(noffH), 0);
    if ((noffH.noffMagic != NOFFMAGIC) && 
		(WordToHost(noffH.noffMagic) == NOFFMAGIC))
    	SwapHeader(&noffH);
    ASSERT(noffH.noffMagic == NOFFMAGIC);
	
    size = noffH.code.size + noffH.initData.size + noffH.uninitData.size ;
    numPages = divRoundUp(size, PageSize) + divRoundUp(UserStackSize,PageSize);
    codeDataSize = divRoundUp(noffH.code.size+noffH.initData.size, PageSize);
	// we need to increase the size
	// to leave room for the stack
    size = numPages * PageSize;
	
    //ASSERT(numPages <= NumPhysPages);		// check we're not trying
	// to run anything too big --
	// at least until we have
	// virtual memory
	
    DEBUG('a', "Initializing address space, num pages %d, size %d\n", 
		  numPages, size);
	// first, set up the translation 
    pageTable = new NewPageEntry[numPages];
	// Allocate physical pages to the user program
	/*phyPageAllocLock->Acquire();
	if ((unsigned int)phyPageMap->NumClear() < numPages) {
		// There is no enough memory
		printf("Error: There is no enough memory\n");
		phyPageAllocLock->Release();
		return;
	}*/
    for (i = 0; i < numPages; i++) {
		pageTable[i].virtualPage = i;	// for now, virtual page # = phys page #
		//int ppn = phyPageMap->Find();	// Find an unused page in physical memory
		//pageTable[i].physicalPage = ppn;
		pageTable[i].valid = FALSE;//TRUE;
		pageTable[i].use = FALSE;
		pageTable[i].dirty = FALSE;
		pageTable[i].readOnly = FALSE;  // if the code segment was entirely on 
		// a separate page, we could set its 
		// pages to be read-only
		if(i < codeDataSize) {
			pageTable[i].location = 1;
			pageTable[i].inFileOffset = noffH.code.inFileAddr+(i*PageSize);
		}else {
			pageTable[i].location = 3;
		}
		pageTable[i].processId = (int) this;
		// zero out the allocated pages
		//bzero(&(machine->mainMemory[pageTable[i].physicalPage*PageSize]), PageSize);
		// populate IPT
		/*ipt[ppn].physicalPage = ppn;
		ipt[ppn].virtualPage = i;
		ipt[ppn].valid = TRUE;
		ipt[ppn].use = FALSE;
		ipt[ppn].dirty = FALSE;
		ipt[ppn].readOnly = FALSE;
		ipt[ppn].location = MEMORY;
		ipt[ppn].processId = this;*/
    }
	//phyPageAllocLock->Release();
	
	// then, copy in the code and data segments into memory
    /*char *buffer;
	if (noffH.code.size > 0) {
        DEBUG('a', "Initializing code segment, at 0x%x, size %d\n", 
			  noffH.code.virtualAddr, noffH.code.size);
		buffer = new char[noffH.code.size];
        executable->ReadAt(buffer,
						   noffH.code.size, noffH.code.inFileAddr);
		for (int j=0; j<noffH.code.size; j++) {
			int physAddr = pageTable[(noffH.code.virtualAddr+j)/PageSize].physicalPage*PageSize
			+(noffH.code.virtualAddr+j)%PageSize;
			machine->mainMemory[physAddr] = buffer[j];
		}
		delete[] buffer;
    }
    if (noffH.initData.size > 0) {
        DEBUG('a', "Initializing data segment, at 0x%x, size %d\n", 
			  noffH.initData.virtualAddr, noffH.initData.size);
        buffer = new char[noffH.initData.size];
        executable->ReadAt(buffer,
						   noffH.initData.size, noffH.initData.inFileAddr);
		for (int j=0; j<noffH.initData.size; j++) {
			int physAddr = pageTable[(noffH.initData.virtualAddr+j)/PageSize].physicalPage*PageSize
			+(noffH.initData.virtualAddr+j)%PageSize;
			machine->mainMemory[physAddr] = buffer[j];
		}
		delete[] buffer;
    }
	if (noffH.uninitData.size > 0) {
        DEBUG('a', "Initializing data segment, at 0x%x, size %d\n", 
			  noffH.uninitData.virtualAddr, noffH.uninitData.size);
        buffer = new char[noffH.uninitData.size];
        executable->ReadAt(buffer,
						   noffH.uninitData.size, noffH.uninitData.inFileAddr);
		for (int j=0; j<noffH.uninitData.size; j++) {
			int physAddr = pageTable[(noffH.uninitData.virtualAddr+j)/PageSize].physicalPage*PageSize
			+(noffH.uninitData.virtualAddr+j)%PageSize;
			machine->mainMemory[physAddr] = buffer[j];
		}
		delete[] buffer;
    }*/
}

//----------------------------------------------------------------------
// AddrSpace::~AddrSpace
//
// 	Dealloate an address space.  release pages, page tables, files
// 	and file tables
//----------------------------------------------------------------------

AddrSpace::~AddrSpace()
{
	DEBUG('w', "Clean up the address space\n");
	phyPageAllocLock->Acquire();
	for (unsigned int i=0; i<numPages; i++) {
		if (pageTable[i].valid) {
			if(ipt[pageTable[i].physicalPage].processId == (int)this) {
				ipt[pageTable[i].physicalPage].valid = FALSE;
				for (int j=0; j<TLBSize; j++) {
					if(machine->tlb[j].valid && machine->tlb[j].physicalPage == pageTable[i].physicalPage) {
						machine->tlb[j].valid = FALSE;
					}
				}
				phyPageMap->Clear(pageTable[i].physicalPage);
			}
		}else if(pageTable[i].location == 2) {
			swapFileLock->Acquire();
			swapFileMap->Clear(pageTable[i].inFileOffset/PageSize);
			swapFileLock->Release();
		}
	}
	phyPageAllocLock->Release();
    delete[] pageTable;
    delete execFile;
}

//----------------------------------------------------------------------
// AddrSpace::InitRegisters
// 	Set the initial values for the user-level register set.
//
// 	We write these directly into the "machine" registers, so
//	that we can immediately jump to user code.  Note that these
//	will be saved/restored into the currentThread->userRegisters
//	when this thread is context switched out.
//----------------------------------------------------------------------

void
AddrSpace::InitRegisters()
{
    int i;

    for (i = 0; i < NumTotalRegs; i++)
	machine->WriteRegister(i, 0);

    // Initial program counter -- must be location of "Start"
    machine->WriteRegister(PCReg, 0);	

    // Need to also tell MIPS where next instruction is, because
    // of branch delay possibility
    machine->WriteRegister(NextPCReg, 4);

   // Set the stack register to the end of the address space, where we
   // allocated the stack; but subtract off a bit, to make sure we don't
   // accidentally reference off the end!
    machine->WriteRegister(StackReg, numPages * PageSize - 16);
    DEBUG('a', "Initializing stack register to %x\n", numPages * PageSize - 16);
}

//----------------------------------------------------------------------
// AddrSpace::SaveState
// 	On a context switch, save any machine state, specific
//	to this address space, that needs saving.
//
//	For now, nothing!
//----------------------------------------------------------------------

void AddrSpace::SaveState() 
{}

//----------------------------------------------------------------------
// AddrSpace::RestoreState
// 	On a context switch, restore the machine state so that
//	this address space can run.
//
//      For now, tell the machine where to find the page table.
//----------------------------------------------------------------------

void AddrSpace::RestoreState() 
{
    //machine->pageTable = pageTable;
    machine->pageTableSize = numPages;
}
/*--------------------------------------------------------------------
 Find the place for the stack of a new thread when Fork() syscall is 
 called.
 Extend the pageTable of the process
 --------------------------------------------------------------------*/
int
AddrSpace::AllocNewStack() {
	unsigned int newStackPageNum = divRoundUp(UserStackSize, PageSize);
	/*phyPageAllocLock->Acquire();
	if ((unsigned int)phyPageMap->NumClear() < newStackPageNum) {
		// memory not enough
		phyPageAllocLock->Release();
		return -1;
	}*/
	NewPageEntry *newPageTable = new NewPageEntry[numPages+newStackPageNum];
	for (unsigned int i=0; i<numPages+newStackPageNum; i++) {
		if (i < numPages) {
			newPageTable[i].virtualPage = pageTable[i].virtualPage;
			newPageTable[i].physicalPage = pageTable[i].physicalPage;
			newPageTable[i].valid = pageTable[i].valid;
			newPageTable[i].use = pageTable[i].use;
			newPageTable[i].dirty = pageTable[i].dirty;
			newPageTable[i].readOnly = pageTable[i].readOnly;
			newPageTable[i].location = pageTable[i].location;
			newPageTable[i].inFileOffset = pageTable[i].inFileOffset;
			newPageTable[i].processId = pageTable[i].processId;
		}else {
			newPageTable[i].virtualPage = i;
			//int ppn = phyPageMap->Find();
			//newPageTable[i].physicalPage = ppn;
			newPageTable[i].valid = FALSE;
			newPageTable[i].use = FALSE;
			newPageTable[i].dirty = FALSE;
			newPageTable[i].readOnly = FALSE;
			newPageTable[i].location = 3;
			newPageTable[i].processId = (int) this;
			//bzero(&(machine->mainMemory[newPageTable[i].physicalPage*PageSize]), PageSize);
			// populate IPT
			/*ipt[ppn].physicalPage = ppn;
			ipt[ppn].virtualPage = i;
			ipt[ppn].valid = TRUE;
			ipt[ppn].use = FALSE;
			ipt[ppn].dirty = FALSE;
			ipt[ppn].readOnly = FALSE;
			ipt[ppn].location = MEMORY;
			ipt[ppn].processId = this;*/
		}
	}
	numPages += newStackPageNum;
	delete[] pageTable;
	pageTable = newPageTable;
	//phyPageAllocLock->Release();
	DEBUG('w', "Increase pages to %d\n", numPages);
	return numPages-1;
}

/*---------------------------------------------------------------------
 * Release the stacks when a thread exits
 * startPageNum: the virtual page number where the stack starts
 ---------------------------------------------------------------------*/
void
AddrSpace::FreeStack(int startPageNum) {
	DEBUG('w', "Free the stack from %d\n", startPageNum);
	unsigned int stackSize = divRoundUp(UserStackSize, PageSize);
	int endPageNum = startPageNum+1-stackSize;
	phyPageAllocLock->Acquire();
	for (int i=endPageNum; i<=startPageNum; i++) {
		if (pageTable[i].valid) {
			pageTable[i].valid = FALSE;
			// clear the IPT
			ipt[pageTable[i].physicalPage].valid = FALSE;
			for (int j=0; j<TLBSize; j++) {
				if(machine->tlb[j].valid && machine->tlb[i].physicalPage == pageTable[i].physicalPage) {
					machine->tlb[j].valid = FALSE;
				}
			}
			phyPageMap->Clear(pageTable[i].physicalPage);
		}else if(pageTable[i].location == 2) {
			swapFileLock->Acquire();
			pageTable[i].location = 0;
			swapFileMap->Clear(pageTable[i].inFileOffset/PageSize);
			swapFileLock->Release();
		}
	}
	phyPageAllocLock->Release();
}

/*-----------------------------------------------------------------
 * Handle the IPT miss, update the IPT
 *----------------------------------------------------------------*/
int
AddrSpace::HandleIPTMiss(int vpn) {
	// Allocate physical pages to the user program
	DEBUG('z', "Handling IPT miss at processId %d, vpn %d\n", (int)this, vpn);
	phyPageAllocLock->Acquire();
	int physPageNum = -1;
	physPageNum = phyPageMap->Find();
	if (physPageNum < 0) {
		// No enough memory, evict page
		physPageNum = HandleMemoryFull();
	}
	if (swapMethod == FIFO) {
		int *tmp = new int;
		*tmp = physPageNum;
		DEBUG('c', "append physical page %d\n", *tmp);
		swapQueue->Append(tmp);
	}
	pageTable[vpn].valid = TRUE;
	pageTable[vpn].physicalPage = physPageNum;
	if (pageTable[vpn].location == 1) {
		// Load it into main memory from executable
		DEBUG('z', "Load page from executable %d\n", pageTable[vpn].inFileOffset);
		execFile->ReadAt(&(machine->mainMemory[physPageNum*PageSize]), PageSize, pageTable[vpn].inFileOffset);
	}else if(pageTable[vpn].location == 2) {
		// Load it into main memory from swap file
		DEBUG('z', "Load page form swap file sector %d\n", pageTable[vpn].inFileOffset/PageSize);
		swapFileLock->Acquire();
		swapFile->ReadAt(&(machine->mainMemory[physPageNum*PageSize]), PageSize, pageTable[vpn].inFileOffset);
		swapFileMap->Clear(pageTable[vpn].inFileOffset/PageSize);
		swapFileLock->Release();
	}
	DEBUG('z', "Update the IPT %d\n", physPageNum);
	ipt[physPageNum].valid = pageTable[vpn].valid;
	ipt[physPageNum].physicalPage = pageTable[vpn].physicalPage;
	ipt[physPageNum].virtualPage = pageTable[vpn].virtualPage;
	ipt[physPageNum].dirty = pageTable[vpn].dirty;
	ipt[physPageNum].use = pageTable[vpn].use;
	ipt[physPageNum].readOnly = pageTable[vpn].readOnly;
	ipt[physPageNum].location = pageTable[vpn].location;
	ipt[physPageNum].inFileOffset = pageTable[vpn].inFileOffset;
	ipt[physPageNum].processId = pageTable[vpn].processId;
	phyPageAllocLock->Release();
	return physPageNum;
}

/*-----------------------------------------------------------------
 * Handle the memory full
 *----------------------------------------------------------------*/
int
AddrSpace::HandleMemoryFull() {
	DEBUG('z', "Handling memory full\n");
	int pageOut;
	if (swapMethod == RANDOM) {
		pageOut = rand()%NumPhysPages;
	}else if(swapMethod == FIFO) {
		int *tmp = (int *)swapQueue->Remove();
		DEBUG('c', "To Evict page %d\n", *tmp);
		pageOut = *tmp;
	}
	//IntStatus old = interrupt->SetLevel(IntOff);
	for (int i=0; i<TLBSize; i++) {
		if (machine->tlb[i].valid && pageOut == machine->tlb[i].physicalPage && ipt[pageOut].processId == (int)this) {
			machine->tlb[i].valid = FALSE;
			if(machine->tlb[i].dirty) {
				// propagate the dirty bit
				ipt[pageOut].dirty = machine->tlb[i].dirty;
				ipt[pageOut].use = machine->tlb[i].use;
			}
		}
	}
	//interrupt->SetLevel(old);
	DEBUG('z', "Evicting the physical page %d\n", pageOut);
	if (ipt[pageOut].dirty) {
		// The page was modified, save it in the swap file
		swapFileLock->Acquire();
		int fileOffset = PageSize*swapFileMap->Find();
		if (fileOffset < 0) {
			printf("Error: The swap file is full.\n");
			return -1;
		}
		DEBUG('z', "Write the physical page %d to swap file sector %d\n", pageOut, fileOffset/PageSize);
		swapFile->WriteAt(&(machine->mainMemory[pageOut*PageSize]), PageSize, fileOffset);
		swapFileLock->Release();
		ipt[pageOut].valid = FALSE;
		ipt[pageOut].location = 2;
		ipt[pageOut].inFileOffset = fileOffset;
	}
	// invalidate the page
	DEBUG('z', "Invalidate the physical page %d of process %d\n", pageOut, ipt[pageOut].processId);
	AddrSpace *addr = (AddrSpace *) ipt[pageOut].processId;
	if(addr != this) {
		DEBUG('x', "current process %d, process %d\n", (int)this, (int)addr);
	}
	addr->pageTable[ipt[pageOut].virtualPage].valid = ipt[pageOut].valid;
	addr->pageTable[ipt[pageOut].virtualPage].dirty = ipt[pageOut].dirty;
	addr->pageTable[ipt[pageOut].virtualPage].use = ipt[pageOut].use;
	addr->pageTable[ipt[pageOut].virtualPage].readOnly = ipt[pageOut].readOnly;
	addr->pageTable[ipt[pageOut].virtualPage].location = ipt[pageOut].location;
	addr->pageTable[ipt[pageOut].virtualPage].inFileOffset = ipt[pageOut].inFileOffset;
	return pageOut;
}


